import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score, recall_score

Xcols = ['eirp',
          'ap_from_ap_0_max_ant_rssi', 'ap_from_ap_0_mean_ant_rssi', 'ap_from_ap_0_sum_ant_rssi',
        # 'ap_from_ap_1_max_ant_rssi', 'ap_from_ap_1_mean_ant_rssi', 'ap_from_ap_1_sum_ant_rssi',
          'sta_from_ap_0_max_ant_rssi', 'sta_from_ap_0_mean_ant_rssi', 'sta_from_ap_0_sum_ant_rssi',
          'sta_from_ap_1_max_ant_rssi', 'sta_from_ap_1_mean_ant_rssi', 'sta_from_ap_1_sum_ant_rssi',
          'sta_to_ap_0_max_ant_rssi', 'sta_to_ap_0_mean_ant_rssi', 'sta_to_ap_0_sum_ant_rssi',
          'sta_to_ap_1_max_ant_rssi', 'sta_to_ap_1_mean_ant_rssi', 'sta_to_ap_1_sum_ant_rssi']

# 读取数据
df = pd.read_csv('training_merge_ReduceRssi.csv')
missing_values = df[Xcols].isnull().any(axis=1)
df = df[~missing_values]
print('data len', len(df))
print('data nss==1 len', len(df[df['nss'] == 1]))

# 分割特征和目标列
X = df[Xcols]
Y = df['nss']

# 创建KNN分类模型
knn = KNeighborsClassifier(n_neighbors=2)

# 训练模型
knn.fit(X, Y)

# 预测测试集
y_pred = knn.predict(X)

# 看看预测nss=1的有多少个
n = 0
for i in y_pred:
    if i == 1:
        n += 1
print('data pred nss==1 len', n)

accuracy = accuracy_score(Y, y_pred)
recall = recall_score(Y, y_pred)
print("acc:", accuracy, 'recall:', recall)

# 跑测试数据
def predTestCsv(filename):
    testData = pd.read_csv(filename)
    testX = testData[Xcols]
    testY = knn.predict(testX)
    print(testY)
    testData['nss'] = pd.Series(testY)
    testData.to_csv(filename)
predTestCsv('test_set_2_2ap_ReduceRssi_pred.csv')
predTestCsv('test_set_2_3ap_ReduceRssi_pred.csv')